{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7e6b66d7-3395-424d-94ee-8bd01d9b724a",
   "metadata": {},
   "source": [
    "# Transformer高级概念示例"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb33d368-2e4b-410a-9d06-30387ed31956",
   "metadata": {},
   "source": [
    "## Transformer各组件变体"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "374cae91-ed93-4a87-b52d-239b0e1bd323",
   "metadata": {},
   "source": [
    "### 自注意力变体示例"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f542eac9-ed57-4101-9472-f756a64e7304",
   "metadata": {},
   "source": [
    "#### 标准多头自注意力"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df9669d6-833b-4873-85ef-6c6a399ed4cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model, num_heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        assert d_model % num_heads == 0\n",
    "        self.d_model = d_model\n",
    "        self.num_heads = num_heads\n",
    "        self.d_head = d_model // num_heads\n",
    "\n",
    "        # 对Q, K, V分别进行线性投影\n",
    "        self.W_q = nn.Linear(d_model, d_model)\n",
    "        self.W_k = nn.Linear(d_model, d_model)\n",
    "        self.W_v = nn.Linear(d_model, d_model)\n",
    "        self.W_o = nn.Linear(d_model, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size, seq_len, d_model = x.size()\n",
    "        \n",
    "        # 计算线性投影\n",
    "        Q = self.W_q(x)  # [batch_size, seq_len, d_model]\n",
    "        K = self.W_k(x)\n",
    "        V = self.W_v(x)\n",
    "        \n",
    "        # 分头\n",
    "        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  # [batch_size, num_heads, seq_len, d_head]\n",
    "        K = K.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)\n",
    "        V = V.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)\n",
    "        \n",
    "        # 缩放点积注意力计算\n",
    "        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)  # [batch_size, num_heads, seq_len, seq_len]\n",
    "        attn = torch.softmax(scores, dim=-1)\n",
    "        context = torch.matmul(attn, V)  # [batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 拼接各头\n",
    "        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)  # [batch_size, seq_len, d_model]\n",
    "        \n",
    "        # 输出投影\n",
    "        output = self.W_o(context)\n",
    "        return output\n",
    "\n",
    "# 使用示例\n",
    "d_model, num_heads, seq_len, batch_size = 512, 8, 64, 2\n",
    "x = torch.rand(batch_size, seq_len, d_model)\n",
    "mha = MultiHeadAttention(d_model, num_heads)\n",
    "output = mha(x)\n",
    "print(\"MHA输出形状:\", output.shape)  # [batch_size, seq_len, d_model]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b3e5014-94d1-42d0-b50f-3569fed80c5b",
   "metadata": {},
   "source": [
    "#### GQA示例\n",
    "\n",
    "GQA 将 $ n_h $ 个查询头分成 $ n_g $ 组，每组共享一组 $ K_j, V_j $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f903e29-c363-4797-a67a-ff91137226cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "\n",
    "class GroupedQueryAttention(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, num_groups):\n",
    "        super(GroupedQueryAttention, self).__init__()\n",
    "        \n",
    "        # 参数校验\n",
    "        # 确保模型维度d_model可被头数num_heads整除，以便分割为多个注意力头\n",
    "        assert d_model % num_heads == 0, \"d_model must be divisible by num_heads\"\n",
    "        # 确保头数num_heads可被分组数num_groups整除，确保每组分配均匀的查询头\n",
    "        assert num_heads % num_groups == 0, \"num_heads must be divisible by num_groups\"\n",
    "        \n",
    "        # 保存模型参数\n",
    "        self.d_model = d_model  # 输入和输出的隐藏状态维度（如 512）\n",
    "        self.num_heads = num_heads  # 注意力头数（如8），决定并行处理的子空间数\n",
    "        self.num_groups = num_groups  # 分组数（如4），决定键和值向量的共享组数\n",
    "        self.d_head = d_model // num_heads  # 每个头的维度（如512/8=64）\n",
    "        self.heads_per_group = num_heads // num_groups  # 每组的查询头数（如8/4=2）\n",
    "\n",
    "        # 定义线性投影层\n",
    "        # 查询（Query）投影矩阵，将输入映射到所有头的维度\n",
    "        self.W_q = nn.Linear(d_model, d_model)  # 形状 [d_model, d_model]\n",
    "        # 键（Key）投影矩阵，仅为num_groups组生成键向量，减少参数量\n",
    "        self.W_k = nn.Linear(d_model, self.d_head * num_groups)  # 形状 [d_model, d_head * num_groups]\n",
    "        # 值（Value）投影矩阵，同样仅为num_groups组生成值向量\n",
    "        self.W_v = nn.Linear(d_model, self.d_head * num_groups)  # 形状 [d_model, d_head * num_groups]\n",
    "        # 输出投影矩阵，将多头输出映射回d_model维度\n",
    "        self.W_o = nn.Linear(d_model, d_model)  # 形状 [d_model, d_model]\n",
    "\n",
    "    def forward(self, x):\n",
    "        # batch_size：批量大小，seq_len：序列长度，d_model：隐藏状态维度\n",
    "        batch_size, seq_len, d_model = x.size()\n",
    "        \n",
    "        # 线性投影，生成查询、键和值向量\n",
    "        Q = self.W_q(x)  # 查询向量，[batch_size, seq_len, d_model]\n",
    "        K = self.W_k(x)  # 键向量，[batch_size, seq_len, d_head * num_groups]\n",
    "        V = self.W_v(x)  # 值向量，[batch_size, seq_len, d_head * num_groups]\n",
    "        \n",
    "        # 分割查询、键和值向量\n",
    "        # 将查询向量Q分割为num_heads个头，形状变为 [batch_size, num_heads, seq_len, d_head]\n",
    "        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)\n",
    "        # 将键向量K分割为num_groups个组，形状变为 [batch_size, num_groups, seq_len, d_head]\n",
    "        K = K.view(batch_size, seq_len, self.num_groups, self.d_head).transpose(1, 2)\n",
    "        # 将值向量V分割为num_groups个组，形状变为 [batch_size, num_groups, seq_len, d_head]\n",
    "        V = V.view(batch_size, seq_len, self.num_groups, self.d_head).transpose(1, 2)\n",
    "        \n",
    "        # 扩展键和值向量以匹配查询头数\n",
    "        # 生成组索引，重复 heads_per_group 次，映射每组到对应的查询头\n",
    "        # 例如：num_groups=4，heads_per_group=2，生成 [0,0,1,1,2,2,3,3]\n",
    "        group_idx = torch.arange(self.num_groups).repeat_interleave(self.heads_per_group)\n",
    "        # 通过索引扩展 K，使每组键向量分配给 heads_per_group 个查询头\n",
    "        K = K[:, group_idx, :, :]  # [batch_size, num_heads, seq_len, d_head]\n",
    "        # 同样扩展 V，使每组值向量分配给 heads_per_group 个查询头\n",
    "        V = V[:, group_idx, :, :]  # [batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 缩放点积注意力\n",
    "        # 计算注意力分数：Q 与 K 的点积，除以 sqrt(d_head) 以稳定梯度\n",
    "        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)  # [batch_size, num_heads, seq_len, seq_len]\n",
    "        # 应用 softmax 归一化，得到注意力权重\n",
    "        attn = torch.softmax(scores, dim=-1)  # [batch_size, num_heads, seq_len, seq_len]\n",
    "        # 使用注意力权重加权值向量，生成上下文向量\n",
    "        context = torch.matmul(attn, V)  # [batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 拼接多头输出\n",
    "        # 将上下文向量转置回 [batch_size, seq_len, num_heads, d_head]，并展平为 [batch_size, seq_len, d_model]\n",
    "        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)\n",
    "        \n",
    "        # 输出投影，将多头输出映射回原始维度\n",
    "        output = self.W_o(context)  # [batch_size, seq_len, d_model]\n",
    "        return output\n",
    "\n",
    "# 示例用法\n",
    "# 设置参数：模型维度 512，8 个注意力头，4 个分组，序列长度 64，批量大小 2\n",
    "d_model, num_heads, num_groups, seq_len, batch_size = 512, 8, 4, 64, 2\n",
    "# 初始化 GQA 模型\n",
    "gqa = GroupedQueryAttention(d_model, num_heads, num_groups)\n",
    "# 生成随机输入张量\n",
    "x = torch.rand(batch_size, seq_len, d_model)\n",
    "# 执行前向传播\n",
    "output = gqa(x)\n",
    "# 打印输出形状，预期为 [batch_size, seq_len, d_model]\n",
    "print(\"GQA的输出形状:\", output.shape)  # [batch_size, seq_len, d_model]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "423274ab-cb6e-4c42-82d3-e38fab1e878e",
   "metadata": {},
   "source": [
    "#### MQA示例\n",
    "\n",
    "MQA 使用单一 $ K, V $，所有 $ n_h $ 个查询头共享"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cac7b63-3719-442b-b986-f6203bd0b128",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "\n",
    "class MultiQueryAttention(nn.Module):\n",
    "    def __init__(self, d_model, num_heads):\n",
    "        super(MultiQueryAttention, self).__init__()\n",
    "        assert d_model % num_heads == 0\n",
    "        self.d_model = d_model\n",
    "        self.num_heads = num_heads\n",
    "        self.d_head = d_model // num_heads\n",
    "\n",
    "        # 线性投影\n",
    "        self.W_q = nn.Linear(d_model, d_model)\n",
    "        self.W_k = nn.Linear(d_model, self.d_head)  # Single K projection\n",
    "        self.W_v = nn.Linear(d_model, self.d_head)  # Single V projection\n",
    "        self.W_o = nn.Linear(d_model, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size, seq_len, d_model = x.size()\n",
    "        \n",
    "        # 线性投影\n",
    "        Q = self.W_q(x)  # [batch_size, seq_len, d_model]\n",
    "        K = self.W_k(x)  # [batch_size, seq_len, d_head]\n",
    "        V = self.W_v(x)  # [batch_size, seq_len, d_head]\n",
    "        \n",
    "        # 将Q分为查询头，并扩展K和V\n",
    "        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  # [batch_size, num_heads, seq_len, d_head]\n",
    "        K = K.unsqueeze(1).expand(-1, self.num_heads, -1, -1)  # [batch_size, num_heads, seq_len, d_head]\n",
    "        V = V.unsqueeze(1).expand(-1, self.num_heads, -1, -1)  # [batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 缩放点积注意力\n",
    "        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)  # [batch_size, num_heads, seq_len, seq_len]\n",
    "        attn = torch.softmax(scores, dim=-1)\n",
    "        context = torch.matmul(attn, V)  # [batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 拼接各头\n",
    "        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)\n",
    "        \n",
    "        # 输出投影\n",
    "        output = self.W_o(context)\n",
    "        return output\n",
    "\n",
    "# Example usage\n",
    "d_model, num_heads, seq_len, batch_size = 512, 8, 64, 2\n",
    "mqa = MultiQueryAttention(d_model, num_heads)\n",
    "x = torch.rand(batch_size, seq_len, d_model)\n",
    "output = mqa(x)\n",
    "print(\"MQA的输出形状:\", output.shape)  # [batch_size, seq_len, d_model]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "168bdaf5-1eb9-4df3-869c-208964b6ac20",
   "metadata": {},
   "source": [
    "#### MLA示例\n",
    "\n",
    "- MLA通过下投影生成低维潜在向量 $ C^{KV}, C^Q \\in \\mathbb{R}^{L \\times d_{\\text{latent}}} $，仅缓存 $ C^{KV} $\n",
    "- 解耦RoPE将潜在向量分为非位置（nope）和位置（rope）部分，简化实现仅对 rope 部分应用旋转\n",
    "  - RoPE指的是Rotary Position Embedding（旋转位置编码），是一种在Transformer模型中广泛使用的相对位置编码方法\n",
    "  - MLA中，RoPE被进一步优化为解耦RoPE（Decoupled RoPE），以适应MLA的低秩压缩机制，确保位置信息在压缩后的潜在空间中能够有效融入注意力计算\n",
    "- 计算复杂度增加（下投影和上投影），内存复杂度为 $ O(L \\cdot d_{\\text{latent}}) $，远低于MHA和GQA\n",
    "\n",
    "**投影层**\n",
    "\n",
    "- 下投影 $ W_{\\text{dkv}}, W_{\\text{dq}} $ 将高维输入压缩到低维潜在空间，减少内存占用。\n",
    "- 上投影 $ W_{\\text{uq}}, W_{\\text{uk}}, W_{\\text{uv}} $ 恢复多头维度，保持表达能力。\n",
    "- 输出投影 $ W_o $ 整合多头输出。\n",
    "\n",
    "**前向传播**\n",
    "\n",
    "- 下投影：将输入 $ x $ 压缩为 $ C_{\\text{kv}}, C_q $，形状为 $ [batch_size, seq_len, d_{\\text{latent}}] $。\n",
    "- 应用 RoPE：对 $ C_{\\text{kv}}, C_q $ 应用解耦 RoPE，融入位置信息。\n",
    "- 上投影：将潜在向量恢复为多头 $ Q, K, V $，形状为 $ [batch_size, num_heads, seq_len, d_head] $。\n",
    "- 注意力计算：执行缩放点积注意力，生成上下文向量。\n",
    "- 输出整合：拼接多头输出并通过 $ W_o $ 投影，返回最终输出和 $ C_{\\text{kv}} $（用于推理缓存）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56b49cb7-0b71-41dc-93a9-93ff7388ade6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "\n",
    "class MultiHeadLatentAttention(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, d_latent, d_rope=64):\n",
    "        super(MultiHeadLatentAttention, self).__init__()\n",
    "        \n",
    "        # 参数校验\n",
    "        # 确保模型维度d_model可被头数num_heads整除，以便分割为多个注意力头\n",
    "        assert d_model % num_heads == 0, \"d_model must be divisible by num_heads\"\n",
    "        # 确保潜在维度d_latent足够大以包含RoPE维度d_rope\n",
    "        assert d_latent >= d_rope, \"d_latent must be at least d_rope\"\n",
    "        # 确保RoPE维度为偶数，因为RoPE需要成对处理维度\n",
    "        assert d_rope % 2 == 0, \"d_rope must be even for RoPE\"\n",
    "        \n",
    "        # 保存模型参数\n",
    "        self.d_model = d_model  # 输入和输出的隐藏状态维度\n",
    "        self.num_heads = num_heads  # 注意力头数\n",
    "        self.d_head = d_model // num_heads  # 每个头的维度\n",
    "        self.d_latent = d_latent  # 潜在空间维度，用于压缩键和值\n",
    "        self.d_rope = d_rope  # 用于RoPE的位置编码维度\n",
    "        self.d_nope = d_latent - d_rope  # 非位置部分（nope）的维度\n",
    "\n",
    "        # 下投影层：将输入从d_model压缩到低维潜在空间d_latent\n",
    "        self.W_dkv = nn.Linear(d_model, d_latent)  # 键和值的共享下投影矩阵\n",
    "        self.W_dq = nn.Linear(d_model, d_latent)   # 查询的下投影矩阵\n",
    "        \n",
    "        # 上投影层：将潜在向量从d_latent恢复到多头维度d_model\n",
    "        self.W_uq = nn.Linear(d_latent, d_model)  # 查询上投影矩阵\n",
    "        self.W_uk = nn.Linear(d_latent, d_model)  # 键上投影矩阵\n",
    "        self.W_uv = nn.Linear(d_latent, d_model)  # 值上投影矩阵\n",
    "        \n",
    "        # 输出投影层：将多头输出映射回d_model维度\n",
    "        self.W_o = nn.Linear(d_model, d_model)\n",
    "        \n",
    "        # 初始化RoPE的频率参数，用于位置编码\n",
    "        self.freqs = self._init_rope_frequencies(d_rope // 2)\n",
    "\n",
    "    def _init_rope_frequencies(self, dim):\n",
    "        # 初始化RoPE的频率参数\n",
    "        # 使用公式 1 / (10000^(2i/d)) 生成频率，dim为d_rope // 2\n",
    "        # torch.arange(0, dim, 1) 生成0到 dim-1 的序列，代表频率索引\n",
    "        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 1).float() / dim))\n",
    "        return inv_freq  # 返回频率张量，形状为 [dim]\n",
    "\n",
    "    def apply_decoupled_rope(self, x, positions):\n",
    "        # 应用解耦 RoPE，将潜在向量分为非位置（nope）和位置（rope）部分，仅对 rope 部分应用旋转\n",
    "        # 输入 x: [batch_size, seq_len, d_latent]，positions: [seq_len]\n",
    "        batch_size, seq_len, _ = x.size()\n",
    "        \n",
    "        # 将潜在向量分割为非位置部分和位置部分\n",
    "        x_nope = x[..., :self.d_nope]  # 非位置部分，[batch_size, seq_len, d_nope]\n",
    "        x_rope = x[..., self.d_nope:]  # 位置部分，[batch_size, seq_len, d_rope]\n",
    "        \n",
    "        # 准备 RoPE 的频率参数\n",
    "        freqs = self.freqs.to(x.device)  # 将频率移到输入设备（如 GPU）\n",
    "        t = positions[:, None].float()  # 位置索引扩展为 [seq_len, 1]\n",
    "        freqs = t * freqs[None, :]  # 计算旋转角度，[seq_len, d_rope // 2]\n",
    "        cos_freq = torch.cos(freqs)  # 余弦值，[seq_len, d_rope // 2]\n",
    "        sin_freq = torch.sin(freqs)  # 正弦值，[seq_len, d_rope // 2]\n",
    "        \n",
    "        # 将 rope 部分分成两半，用于成对旋转\n",
    "        x1 = x_rope[..., :self.d_rope // 2]  # 前半部分，[batch_size, seq_len, d_rope // 2]\n",
    "        x2 = x_rope[..., self.d_rope // 2:]  # 后半部分，[batch_size, seq_len, d_rope // 2]\n",
    "        \n",
    "        # 应用 RoPE 旋转\n",
    "        # 旋转公式：(x1, x2) -> (x1 * cos - x2 * sin, x1 * sin + x2 * cos)\n",
    "        x_rope_rotated = torch.cat([\n",
    "            x1 * cos_freq - x2 * sin_freq,  # 旋转后的前半部分\n",
    "            x1 * sin_freq + x2 * cos_freq   # 旋转后的后半部分\n",
    "        ], dim=-1)  # 拼接，[batch_size, seq_len, d_rope]\n",
    "        \n",
    "        # 拼接非位置部分和旋转后的位置部分\n",
    "        return torch.cat([x_nope, x_rope_rotated], dim=-1)  # [batch_size, seq_len, d_latent]\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 前向传播，输入 x: [batch_size, seq_len, d_model]\n",
    "        batch_size, seq_len, d_model = x.size()\n",
    "        \n",
    "        # 下投影：将输入压缩到低维潜在空间\n",
    "        C_kv = self.W_dkv(x)  # 键和值潜在向量，[batch_size, seq_len, d_latent]\n",
    "        C_q = self.W_dq(x)    # 查询潜在向量，[batch_size, seq_len, d_latent]\n",
    "        \n",
    "        # 应用解耦 RoPE，为潜在向量添加位置信息\n",
    "        positions = torch.arange(seq_len, device=x.device)  # 位置索引 [0, 1, ..., seq_len-1]\n",
    "        C_kv = self.apply_decoupled_rope(C_kv, positions)  # 对键和值潜在向量应用 RoPE\n",
    "        C_q = self.apply_decoupled_rope(C_q, positions)    # 对查询潜在向量应用 RoPE\n",
    "        \n",
    "        # 上投影：将潜在向量恢复到多头维度\n",
    "        Q = self.W_uq(C_q).view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  # 查询，[batch_size, num_heads, seq_len, d_head]\n",
    "        K = self.W_uk(C_kv).view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  # 键，[batch_size, num_heads, seq_len, d_head]\n",
    "        V = self.W_uv(C_kv).view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  # 值，[batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 缩放点积注意力\n",
    "        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)  # 注意力分数，[batch_size, num_heads, seq_len, seq_len]\n",
    "        attn = torch.softmax(scores, dim=-1)  # 注意力权重，归一化\n",
    "        context = torch.matmul(attn, V)  # 上下文向量，[batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 拼接多头输出\n",
    "        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)  # [batch_size, seq_len, d_model]\n",
    "        \n",
    "        # 输出投影\n",
    "        output = self.W_o(context)  # 最终输出，[batch_size, seq_len, d_model]\n",
    "        return output, C_kv  # 返回输出和键-值潜在向量（用于缓存）\n",
    "\n",
    "# 示例用法\n",
    "d_model, num_heads, d_latent, d_rope, seq_len, batch_size = 512, 8, 128, 64, 64, 2\n",
    "mla = MultiHeadLatentAttention(d_model, num_heads, d_latent, d_rope)\n",
    "x = torch.rand(batch_size, seq_len, d_model)\n",
    "output, C_kv = mla(x)\n",
    "print(\"MLA的输出形状:\", output.shape)  # [batch_size, seq_len, d_model]\n",
    "print(\"MLA KV cache的形状:\", C_kv.shape)  # [batch_size, seq_len, d_latent]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e22ce0e5-67af-40f3-aeec-5466e4a8d595",
   "metadata": {},
   "source": [
    "### 局部注意力机制"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fea0e301-71bc-44a7-b275-2b479e18281e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "\n",
    "class LocalAttention(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, window_size):\n",
    "        super(LocalAttention, self).__init__()\n",
    "        \n",
    "        # 参数校验\n",
    "        # 确保模型维度d_model可被头数num_heads整除，以便分割为多个注意力头\n",
    "        assert d_model % num_heads == 0, \"d_model must be divisible by num_heads\"\n",
    "        \n",
    "        # 保存模型参数\n",
    "        self.d_model = d_model  # 输入和输出的隐藏状态维度（如 512）\n",
    "        self.num_heads = num_heads  # 注意力头数（如8），决定并行处理的子空间数\n",
    "        self.d_head = d_model // num_heads  # 每个头的维度（如 512/8=64）\n",
    "        self.window_size = window_size  # 滑动窗口的半宽（左右各 window_size 个 token）\n",
    "\n",
    "        # 定义线性投影层\n",
    "        # 查询（Query）投影矩阵，将输入映射到所有头的维度\n",
    "        self.W_q = nn.Linear(d_model, d_model)  # 形状 [d_model, d_model]\n",
    "        # 键（Key）投影矩阵，生成所有头的键向量\n",
    "        self.W_k = nn.Linear(d_model, d_model)  # 形状 [d_model, d_model]\n",
    "        # 值（Value）投影矩阵，生成所有头的值向量\n",
    "        self.W_v = nn.Linear(d_model, d_model)  # 形状 [d_model, d_model]\n",
    "        # 输出投影矩阵，将多头输出映射回d_model维度\n",
    "        self.W_o = nn.Linear(d_model, d_model)  # 形状 [d_model, d_model]\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 前向传播，输入 x: [batch_size, seq_len, d_model]\n",
    "        # batch_size：批量大小，seq_len：序列长度，d_model：隐藏状态维度\n",
    "        batch_size, seq_len, d_model = x.size()\n",
    "        \n",
    "        # 线性投影，生成查询、键和值向量\n",
    "        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  # 查询，[batch_size, num_heads, seq_len, d_head]\n",
    "        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  # 键，[batch_size, num_heads, seq_len, d_head]\n",
    "        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  # 值，[batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 初始化上下文向量，用于存储注意力计算结果\n",
    "        context = torch.zeros_like(Q)  # [batch_size, num_heads, seq_len, d_head]\n",
    "        \n",
    "        # 滑动窗口局部注意力\n",
    "        for i in range(seq_len):\n",
    "            # 定义当前token i的窗口范围\n",
    "            # 窗口覆盖 [i-window_size, i+window_size]，并处理边界情况\n",
    "            start = max(0, i - self.window_size)  # 窗口起始位置，最小为0\n",
    "            end = min(seq_len, i + self.window_size + 1)  # 窗口结束位置，最大为seq_len\n",
    "            \n",
    "            # 提取局部窗口内的键和值向量\n",
    "            K_local = K[:, :, start:end, :]  # [batch_size, num_heads, window_size, d_head]\n",
    "            V_local = V[:, :, start:end, :]  # [batch_size, num_heads, window_size, d_head]\n",
    "            \n",
    "            # 获取当前token i的查询向量\n",
    "            q_i = Q[:, :, i:i+1, :]  # [batch_size, num_heads, 1, d_head]\n",
    "            \n",
    "            # 计算局部注意力分数\n",
    "            # q_i与K_local点积，除以sqrt(d_head)稳定梯度\n",
    "            scores = torch.matmul(q_i, K_local.transpose(-2, -1)) / math.sqrt(self.d_head)  # [batch_size, num_heads, 1, window_size]\n",
    "            \n",
    "            # 应用softmax归一化，得到局部注意力权重\n",
    "            attn = torch.softmax(scores, dim=-1)  # [batch_size, num_heads, 1, window_size]\n",
    "            \n",
    "            # 使用注意力权重加权局部值向量，生成上下文向量\n",
    "            context[:, :, i:i+1, :] = torch.matmul(attn, V_local)  # [batch_size, num_heads, 1, d_head]\n",
    "        \n",
    "        # 拼接多头输出\n",
    "        # 将上下文向量转置回 [batch_size, seq_len, num_heads, d_head]，并展平为 [batch_size, seq_len, d_model]\n",
    "        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)\n",
    "        \n",
    "        # 输出投影，将多头输出映射回原始维度\n",
    "        output = self.W_o(context)  # [batch_size, seq_len, d_model]\n",
    "        return output\n",
    "\n",
    "# 示例用法\n",
    "# 设置参数：模型维度512，8个注意力头，窗口半宽16，序列长度64，批量大小2\n",
    "d_model, num_heads, window_size, seq_len, batch_size = 512, 8, 16, 64, 2\n",
    "# 初始化局部注意力模型\n",
    "local_attn = LocalAttention(d_model, num_heads, window_size)\n",
    "# 生成随机输入张量\n",
    "x = torch.rand(batch_size, seq_len, d_model)\n",
    "# 执行前向传播\n",
    "output = local_attn(x)\n",
    "# 打印输出形状，预期为[batch_size, seq_len, d_model]\n",
    "print(\"Local Attention的输出形状:\", output.shape)  # [batch_size, seq_len, d_model]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c3b0d7-f11b-41c5-8f43-acd9c9b62a29",
   "metadata": {},
   "source": [
    "## 层归一化和残差连接"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61a57e26-a962-4f56-9ce3-7ec58abfe382",
   "metadata": {},
   "source": [
    "### 层归一化"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a43e412b-4a28-499e-a58c-caf9c8b0f42d",
   "metadata": {},
   "source": [
    "### 残差连接"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e1e33dd-9c25-4240-abbf-1ac480ed8968",
   "metadata": {},
   "source": [
    "#### 双残差连接\n",
    "\n",
    "将输入特征分成两条路径分别进行处理，一直负责网络内部计算，一个负责保证梯度稳定。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85cf9806-0a25-4431-b759-530f5393804a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# 定义Decoder-Only Transformer的层，融入标准的双残差ResiDual设计\n",
    "class ResiDualTransformerBlock(nn.Module):\n",
    "    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):\n",
    "        \"\"\"\n",
    "        初始化ResiDualTransformerBlock模块。\n",
    "        \n",
    "        参数:\n",
    "        - d_model: 模型的嵌入维度（hidden size）。\n",
    "        - n_heads: 自注意力机制中的注意力头数。\n",
    "        - d_ff: 前馈网络（FFN）的中间维度。\n",
    "        - dropout: dropout率，用于防止过拟合。\n",
    "        \"\"\"\n",
    "        super(ResiDualTransformerBlock, self).__init__()\n",
    "        # 自注意力层（masked multi-head attention），对应f(Attn(·))\n",
    "        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)\n",
    "        # 前馈网络（FFN），对应g(·)，通常是两个线性层加激活函数\n",
    "        self.ffn = nn.Sequential(\n",
    "            nn.Linear(d_model, d_ff),  # 第一层线性变换\n",
    "            nn.GELU(),                 # GELU激活函数\n",
    "            nn.Linear(d_ff, d_model),  # 第二层线性变换\n",
    "            nn.Dropout(dropout)        # dropout层\n",
    "        )\n",
    "        # 层归一化（Layer Normalization）模块，用于非线性变换\n",
    "        self.ln1 = nn.LayerNorm(d_model)  # 第一次LN，用于输入归一化\n",
    "        self.ln2 = nn.LayerNorm(d_model)  # 第二次LN，用于第一次残差后归一化\n",
    "        self.ln3 = nn.LayerNorm(d_model)  # 第三次LN，用于第二次残差后归一化\n",
    "        # dropout层，用于注意力输出\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, attn_mask=None):\n",
    "        \"\"\"\n",
    "        前向传播函数，实现标准的双残差ResiDual流程。\n",
    "        \n",
    "        参数:\n",
    "        - x: 输入张量，形状为(batch_size, seq_len, d_model)。\n",
    "        - attn_mask: 可选的注意力掩码，用于Decoder-Only的因果注意力（causal mask）。\n",
    "        \n",
    "        流程:\n",
    "        1. 对输入x进行层归一化，得到x_ln。\n",
    "        2. 计算自注意力（masked），并与x_ln相加，形成第一次残差连接（y_local1）。\n",
    "        3. 对y_local1进行层归一化，得到y_local1_ln。\n",
    "        4. 通过FFN处理y_local1_ln，并与y_local1相加，形成第二次残差连接（y_local2）。\n",
    "        5. 对y_local2进行层归一化，得到中间输出。\n",
    "        6. 最后将原始输入x与中间输出相加，形成全局残差连接，得到最终输出y。\n",
    "        \"\"\"\n",
    "        # 步骤1: 输入归一化\n",
    "        x_ln = self.ln1(x)\n",
    "        \n",
    "        # 步骤2: 自注意力计算（f(Attn(x_ln))），注意要使用attn_mask实现Decoder-Only的因果注意力\n",
    "        attn_output, _ = self.self_attn(x_ln, x_ln, x_ln, attn_mask=attn_mask)\n",
    "        attn_output = self.dropout(attn_output)\n",
    "        \n",
    "        # 第一次残差连接\n",
    "        y_local1 = x_ln + attn_output\n",
    "        \n",
    "        # 步骤3: 对y_local1进行归一化\n",
    "        y_local1_ln = self.ln2(y_local1)\n",
    "        \n",
    "        # 步骤4: 前馈网络计算（g(FFN(y_local1_ln))）\n",
    "        ffn_output = self.ffn(y_local1_ln)\n",
    "        \n",
    "        # 第二次残差连接\n",
    "        y_local2 = y_local1 + ffn_output\n",
    "        \n",
    "        # 步骤5: 对y_local2进行归一化，得到中间输出\n",
    "        intermediate = self.ln3(y_local2)\n",
    "        \n",
    "        # 步骤6: 全局残差连接: y = x + intermediate\n",
    "        y = x + intermediate\n",
    "        \n",
    "        return y\n",
    "\n",
    "# 示例使用: 创建一个ResiDualTransformerBlock实例，并进行前向传播测试\n",
    "if __name__ == \"__main__\":\n",
    "    # 定义参数\n",
    "    batch_size = 2\n",
    "    seq_len = 10\n",
    "    d_model = 512\n",
    "    n_heads = 8\n",
    "    d_ff = 2048\n",
    "    \n",
    "    # 创建模块\n",
    "    block = ResiDualTransformerBlock(d_model, n_heads, d_ff)\n",
    "    \n",
    "    # 生成随机输入\n",
    "    x = torch.randn(batch_size, seq_len, d_model)\n",
    "    \n",
    "    # 生成因果注意力掩码（Decoder-Only）\n",
    "    attn_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1) * float('-inf')\n",
    "    \n",
    "    # 前向传播\n",
    "    output = block(x, attn_mask=attn_mask)\n",
    "    \n",
    "    # 打印输出形状以验证\n",
    "    print(\"输出形状:\", output.shape)  # 应为 (batch_size, seq_len, d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2eeddf2d-badb-4bce-a5cc-a827c3eef09e",
   "metadata": {},
   "source": [
    "#### 并行结构\n",
    "\n",
    "将MSA子层和FFN子层并行放置，让它们同时处理输入。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f4d5f0-5b6c-4c58-b274-2bb0318f6e44",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# 定义一个带有并行结构的Transformer层\n",
    "class ParallelTransformerBlock(nn.Module):\n",
    "    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):\n",
    "        \"\"\"\n",
    "        初始化ParallelTransformerBlock模块。\n",
    "        \n",
    "        参数:\n",
    "        - d_model: 模型的嵌入维度（hidden size）。\n",
    "        - n_heads: 自注意力机制中的注意力头数。\n",
    "        - d_ff: 前馈网络（FFN）的中间维度。\n",
    "        - dropout: dropout率，用于防止过拟合。\n",
    "        \"\"\"\n",
    "        super(ParallelTransformerBlock, self).__init__()\n",
    "        \n",
    "        # 层归一化（Layer Normalization），用于并行分支的输入归一化\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "        \n",
    "        # 自注意力机制（Multi-Head Self-Attention）\n",
    "        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)\n",
    "        \n",
    "        # 前馈网络（FFN），包含两个线性层和激活函数\n",
    "        self.ffn = nn.Sequential(\n",
    "            nn.Linear(d_model, d_ff),  # 第一层线性变换\n",
    "            nn.GELU(),                 # GELU激活函数\n",
    "            nn.Linear(d_ff, d_model),  # 第二层线性变换\n",
    "            nn.Dropout(dropout)        # dropout层\n",
    "        )\n",
    "        \n",
    "        # dropout层，用于并行输出的正则化\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, attn_mask=None):\n",
    "        \"\"\"\n",
    "        前向传播函数，实现并行自注意力和FFN结构。\n",
    "        \n",
    "        参数:\n",
    "        - x: 输入张量，形状为(batch_size, seq_len, d_model)。\n",
    "        - attn_mask: 可选的注意力掩码，用于Decoder-Only的因果注意力（causal mask）。\n",
    "        \n",
    "        流程:\n",
    "        1. 对输入x进行层归一化，得到x_ln。\n",
    "        2. 并行计算自注意力输出（attn_out）和FFN输出（ffn_out）。\n",
    "        3. 将注意力输出和FFN输出与归一化后的输入x_ln相加。\n",
    "        4. 应用dropout正则化。\n",
    "        5. 全局残差连接，将结果与原始输入x相加。\n",
    "        6. 返回最终输出。\n",
    "        \"\"\"\n",
    "        # 步骤1: 输入归一化\n",
    "        x_ln = self.ln(x)\n",
    "        \n",
    "        # 步骤2: 并行计算自注意力输出和FFN输出\n",
    "        # 自注意力分支\n",
    "        attn_out, _ = self.self_attn(x_ln, x_ln, x_ln, attn_mask=attn_mask)\n",
    "        attn_out = self.dropout(attn_out)\n",
    "        \n",
    "        # FFN分支\n",
    "        ffn_out = self.ffn(x_ln)\n",
    "        ffn_out = self.dropout(ffn_out)\n",
    "        \n",
    "        # 步骤3: 并行融合，结合自注意力输出和FFN输出\n",
    "        parallel_out = x_ln + attn_out + ffn_out\n",
    "        \n",
    "        # 步骤4: 应用dropout正则化\n",
    "        parallel_out = self.dropout(parallel_out)\n",
    "        \n",
    "        # 步骤5: 全局残差连接\n",
    "        y = x + parallel_out\n",
    "        \n",
    "        return y\n",
    "\n",
    "# 示例使用: 创建一个ParallelTransformerBlock实例，并进行前向传播测试\n",
    "if __name__ == \"__main__\":\n",
    "    # 定义参数\n",
    "    batch_size = 2\n",
    "    seq_len = 10\n",
    "    d_model = 512\n",
    "    n_heads = 8\n",
    "    d_ff = 2048\n",
    "    \n",
    "    # 创建模块\n",
    "    block = ParallelTransformerBlock(d_model, n_heads, d_ff)\n",
    "    \n",
    "    # 生成随机输入\n",
    "    x = torch.randn(batch_size, seq_len, d_model)\n",
    "    \n",
    "    # 生成因果注意力掩码（适用于Decoder-Only）\n",
    "    attn_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1) * float('-inf')\n",
    "    \n",
    "    # 前向传播\n",
    "    output = block(x, attn_mask=attn_mask)\n",
    "    \n",
    "    # 打印输出形状以验证\n",
    "    print(\"输出形状:\", output.shape)  # 应为 (batch_size, seq_len, d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be488aea-c3fb-463f-b922-e141fef6a410",
   "metadata": {},
   "source": [
    "#### 随机深度\n",
    "\n",
    "对于残差连接𝑦=𝑥+𝐹(𝑥)，在训练期间以一定概率随机丢弃残差分支中的𝐹(𝑥)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dd3cc01-635d-42ef-bca9-b3585f91052a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class StochasticDepth(nn.Module):\n",
    "    def __init__(self, drop_prob: float):\n",
    "        \"\"\"\n",
    "        drop_prob: 丢弃概率 (0 <= drop_prob < 1)\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.drop_prob = drop_prob\n",
    "\n",
    "    def forward(self, x, residual):\n",
    "        \"\"\"\n",
    "        x: 残差主分支 (输入)\n",
    "        residual: 子层输出 f(x)\n",
    "        \"\"\"\n",
    "        if not self.training or self.drop_prob == 0.0:\n",
    "            # 推理时（eval模式）或不丢弃 → 标准残差\n",
    "            return x + residual\n",
    "\n",
    "        # 训练时：生成 Bernoulli 掩码\n",
    "        keep_prob = 1 - self.drop_prob\n",
    "        mask = torch.empty(x.shape[0], 1, 1, device=x.device).bernoulli_(keep_prob)\n",
    "        # 保持期望不变，需要缩放 residual\n",
    "        residual = residual / keep_prob\n",
    "        return x + residual * mask\n",
    "\n",
    "\n",
    "# ===== 使用示例：Transformer Block =====\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, d_model, nhead, dim_ff, drop_prob=0.1):\n",
    "        super().__init__()\n",
    "        self.ln1 = nn.LayerNorm(d_model)\n",
    "        self.ln2 = nn.LayerNorm(d_model)\n",
    "        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)\n",
    "        self.ffn = nn.Sequential(\n",
    "            nn.Linear(d_model, dim_ff),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(dim_ff, d_model)\n",
    "        )\n",
    "        # 在残差连接中应用随机深度\n",
    "        self.sd1 = StochasticDepth(drop_prob)\n",
    "        self.sd2 = StochasticDepth(drop_prob)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Self-Attention + 残差（带随机深度）\n",
    "        attn_out, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x))\n",
    "        x = self.sd1(x, attn_out)\n",
    "\n",
    "        # FFN + 残差（带随机深度）\n",
    "        ffn_out = self.ffn(self.ln2(x))\n",
    "        x = self.sd2(x, ffn_out)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "# ===== 测试运行 =====\n",
    "if __name__ == \"__main__\":\n",
    "    torch.manual_seed(42)\n",
    "\n",
    "    batch_size, seq_len, d_model = 2, 10, 512\n",
    "    x = torch.randn(batch_size, seq_len, d_model)\n",
    "\n",
    "    block = TransformerBlock(d_model=512, nhead=4, dim_ff=64, drop_prob=0.2)\n",
    "    block.train()  # 训练模式：随机丢弃\n",
    "\n",
    "    out = block(x)\n",
    "    print(\"Output shape (train):\", out.shape)\n",
    "\n",
    "    block.eval()  # 推理模式：不丢弃\n",
    "    out_eval = block(x)\n",
    "    print(\"Output shape (eval):\", out_eval.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa2494bb-150a-4405-8a03-39347bed5022",
   "metadata": {},
   "source": [
    "## Transformer MoE模型示例"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53fc1060-45de-4ba1-94a9-0a6974d2cbcb",
   "metadata": {},
   "source": [
    "### 简单MoE模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67331ad9-3943-4023-a039-3f674c68c07b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "\n",
    "# 多头注意力机制\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model, num_heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        assert d_model % num_heads == 0, \"d_model must be divisible by num_heads\"\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        self.num_heads = num_heads\n",
    "        self.d_k = d_model // num_heads\n",
    "        \n",
    "        self.W_q = nn.Linear(d_model, d_model)\n",
    "        self.W_k = nn.Linear(d_model, d_model)\n",
    "        self.W_v = nn.Linear(d_model, d_model)\n",
    "        self.W_o = nn.Linear(d_model, d_model)\n",
    "        \n",
    "    def scaled_dot_product_attention(self, Q, K, V, mask=None):\n",
    "        batch_size = Q.size(0)\n",
    "        QK = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)\n",
    "        \n",
    "        if mask is not None:\n",
    "            QK = QK.masked_fill(mask == 0, float('-inf'))\n",
    "        \n",
    "        attention_scores = F.softmax(QK, dim=-1)\n",
    "        output = torch.matmul(attention_scores, V)\n",
    "        return output, attention_scores\n",
    "    \n",
    "    def forward(self, x, mask=None):\n",
    "        batch_size = x.size(0)\n",
    "        \n",
    "        Q = self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)\n",
    "        K = self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)\n",
    "        V = self.W_v(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)\n",
    "        \n",
    "        output, attention_scores = self.scaled_dot_product_attention(Q, K, V, mask)\n",
    "        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)\n",
    "        output = self.W_o(output)\n",
    "        return output, attention_scores\n",
    "\n",
    "# 混合专家系统 (MOE)\n",
    "class MoEFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, d_ff, num_experts, top_k=2):\n",
    "        super(MoEFeedForward, self).__init__()\n",
    "        self.num_experts = num_experts\n",
    "        self.top_k = top_k\n",
    "        \n",
    "        # 门控网络\n",
    "        self.gate = nn.Linear(d_model, num_experts)\n",
    "        \n",
    "        # 专家网络\n",
    "        self.experts = nn.ModuleList([\n",
    "            nn.Sequential(\n",
    "                nn.Linear(d_model, d_ff),\n",
    "                nn.ReLU(),\n",
    "                nn.Linear(d_ff, d_model)\n",
    "            ) for _ in range(num_experts)\n",
    "        ])\n",
    "        \n",
    "    def forward(self, x):\n",
    "        batch_size, seq_len, d_model = x.size()\n",
    "        \n",
    "        # 门控网络输出专家选择概率\n",
    "        gate_scores = F.softmax(self.gate(x), dim=-1)  # [batch_size, seq_len, num_experts]\n",
    "        \n",
    "        # 选择 top-k 专家\n",
    "        top_k_scores, top_k_indices = gate_scores.topk(self.top_k, dim=-1)  # [batch_size, seq_len, top_k]\n",
    "        \n",
    "        # 初始化输出\n",
    "        output = torch.zeros_like(x)\n",
    "        \n",
    "        for i in range(self.top_k):\n",
    "            expert_idx = top_k_indices[:, :, i]  # [batch_size, seq_len]\n",
    "            expert_scores = top_k_scores[:, :, i].unsqueeze(-1)  # [batch_size, seq_len, 1]\n",
    "            \n",
    "            # 为每个样本选择对应的专家输出\n",
    "            expert_output = torch.zeros_like(x)\n",
    "            for j in range(self.num_experts):\n",
    "                mask = (expert_idx == j).unsqueeze(-1).float()\n",
    "                expert_output += mask * self.experts[j](x)\n",
    "            \n",
    "            output += expert_output * expert_scores\n",
    "        \n",
    "        return output\n",
    "\n",
    "# Transformer Decoder 层\n",
    "class TransformerDecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, d_ff, num_experts, dropout=0.1):\n",
    "        super(TransformerDecoderLayer, self).__init__()\n",
    "        self.self_attention = MultiHeadAttention(d_model, num_heads)\n",
    "        self.norm1 = nn.LayerNorm(d_model)\n",
    "        self.moe = MoEFeedForward(d_model, d_ff, num_experts)\n",
    "        self.norm2 = nn.LayerNorm(d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x, mask=None):\n",
    "        # Pre-LayerNorm\n",
    "        x_norm = self.norm1(x)\n",
    "        attn_output, attn_scores = self.self_attention(x_norm, mask)\n",
    "        x = x + self.dropout(attn_output)\n",
    "        \n",
    "        # Pre-LayerNorm for MoE\n",
    "        x_norm = self.norm2(x)\n",
    "        moe_output = self.moe(x_norm)\n",
    "        x = x + self.dropout(moe_output)\n",
    "        \n",
    "        return x, attn_scores\n",
    "\n",
    "# 完整的 Transformer Decoder-Only 模型\n",
    "class TransformerDecoderOnly(nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, num_experts, max_seq_len, dropout=0.1):\n",
    "        super(TransformerDecoderOnly, self).__init__()\n",
    "        self.d_model = d_model\n",
    "        self.embedding = nn.Embedding(vocab_size, d_model)\n",
    "        self.pos_encoding = self.create_positional_encoding(max_seq_len, d_model)\n",
    "        self.layers = nn.ModuleList([\n",
    "            TransformerDecoderLayer(d_model, num_heads, d_ff, num_experts, dropout)\n",
    "            for _ in range(num_layers)\n",
    "        ])\n",
    "        self.fc_out = nn.Linear(d_model, vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def create_positional_encoding(self, max_seq_len, d_model):\n",
    "        pe = torch.zeros(max_seq_len, d_model)\n",
    "        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0)  # [1, max_seq_len, d_model]\n",
    "        return pe\n",
    "    \n",
    "    def create_causal_mask(self, seq_len):\n",
    "        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()\n",
    "        return ~mask  # [seq_len, seq_len]\n",
    "    \n",
    "    def forward(self, x):\n",
    "        batch_size, seq_len = x.size()\n",
    "        mask = self.create_causal_mask(seq_len).to(x.device)\n",
    "        \n",
    "        # 嵌入和位置编码\n",
    "        x = self.embedding(x) * math.sqrt(self.d_model)\n",
    "        x = x + self.pos_encoding[:, :seq_len, :].to(x.device)\n",
    "        x = self.dropout(x)\n",
    "        \n",
    "        # 通过所有解码器层\n",
    "        attention_scores = []\n",
    "        for layer in self.layers:\n",
    "            x, attn = layer(x, mask)\n",
    "            attention_scores.append(attn)\n",
    "        \n",
    "        # 输出层\n",
    "        output = self.fc_out(x)\n",
    "        return output, attention_scores\n",
    "\n",
    "# 示例使用\n",
    "def main():\n",
    "    # 模型参数\n",
    "    vocab_size = 10000\n",
    "    d_model = 512\n",
    "    num_heads = 8\n",
    "    d_ff = 2048\n",
    "    num_layers = 6\n",
    "    num_experts = 4\n",
    "    max_seq_len = 100\n",
    "    dropout = 0.1\n",
    "    \n",
    "    # 初始化模型\n",
    "    model = TransformerDecoderOnly(\n",
    "        vocab_size, d_model, num_heads, d_ff, num_layers, num_experts, max_seq_len, dropout\n",
    "    )\n",
    "    \n",
    "    # 模拟输入\n",
    "    batch_size = 32\n",
    "    seq_len = 50\n",
    "    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))\n",
    "    \n",
    "    # 前向传播\n",
    "    output, attention_scores = model(input_ids)\n",
    "    print(f\"Output shape: {output.shape}\")  # [batch_size, seq_len, vocab_size]\n",
    "    print(f\"Number of attention score tensors: {len(attention_scores)}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
